package com.analysis.client;

import com.analysis.dto.AiCallResponse;
import com.analysis.utils.ParseTextUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import java.util.*;

/**
 * KIMI AI 高可用客户端
 */
@Component
public class KimiAIChatClient {

    private static final Logger logger = LoggerFactory.getLogger(KimiAIChatClient.class);
    private static final Logger aiModelLogger = LoggerFactory.getLogger("com.analysis.client");

    @Autowired
    private RestTemplate restTemplate;

    @Value("${kimi.chat.api-key:}")
    private String apiKey;

    @Value("${kimi.chat.base-url:https://api.moonshot.cn/v1/chat/completions}")
    private String baseUrl;

    @Value("${kimi.chat.model:kimi-k2-0905-preview}")
    private String model;

    @Value("${kimi.chat.max-retries:2}")
    private int maxRetries;

    @Value("${kimi.chat.base-delay-ms:1000}")
    private long baseDelayMs;

    @Value("${kimi.chat.max-delay-ms:10000}")
    private long maxDelayMs;

    private final Random random = new Random();

    /**
     * 联网搜索问答（两阶段：调用内置$web_search工具 → 注入tool结果再次请求）
     * - 复用全局RestTemplate超时（5分钟）
     * - 带指数退避重试
     */
    public AiCallResponse webSearch(String query) {
        Exception lastException = null;
        for (int attempt = 1; attempt <= Math.max(1, maxRetries); attempt++) {
            try {
                aiModelLogger.info("第{}次尝试调用 KIMI API(联网) - 问题: {}", attempt, query);

                // 第一次：触发tool_calls
                ChatRequest first = buildFirstRequest(query);
                ChatResponse firstResp = callKimi(first);

                if (firstResp == null || firstResp.choices == null || firstResp.choices.isEmpty()) {
                    throw new RuntimeException("KIMI 首次响应为空或无choices");
                }

                Choice firstChoice = firstResp.choices.get(0);
                if (firstChoice.message == null || firstChoice.message.tool_calls == null || firstChoice.message.tool_calls.isEmpty()) {
                    // 没有触发工具，直接返回内容（退化处理）
                    return buildAiResponse(firstChoice.message == null ? null : firstChoice.message.content);
                }

                ToolCall toolCall = firstChoice.message.tool_calls.get(0);
                if (toolCall.function == null || toolCall.function.arguments == null) {
                    throw new RuntimeException("KIMI tool_calls 缺少function.arguments");
                }

                // 第二次：携带assistant的tool_calls与tool角色的结果
                ChatRequest second = buildSecondRequest(query, firstChoice.message, toolCall);
                ChatResponse secondResp = callKimi(second);

                if (secondResp != null && secondResp.choices != null && !secondResp.choices.isEmpty()) {
                    String content = Optional.ofNullable(secondResp.choices.get(0))
                            .map(c -> c.message)
                            .map(m -> m.content)
                            .orElse(null);
                    return buildAiResponse(content);
                }

                throw new RuntimeException("KIMI 二次响应为空或无choices");
            } catch (Exception e) {
                lastException = e;
                aiModelLogger.warn("第{}次 KIMI API(联网) 调用失败 - 错误: {}", attempt, e.getMessage());
                if (attempt < Math.max(1, maxRetries)) {
                    long delay = calculateBackoffDelay(attempt);
                    aiModelLogger.info("等待 {}ms 后进行第{}次重试", delay, attempt + 1);
                    try {
                        Thread.sleep(delay);
                    } catch (InterruptedException ie) {
                        Thread.currentThread().interrupt();
                        throw new RuntimeException("KIMI API 重试被中断", ie);
                    }
                }
            }
        }
        aiModelLogger.error("KIMI API(联网) 调用最终失败 - 已重试{}次, 最后错误: {}",
                Math.max(1, maxRetries), lastException != null ? lastException.getMessage() : "未知错误");
        return buildAiResponse(null);
    }

    private ChatRequest buildFirstRequest(String query) {
        ChatRequest req = new ChatRequest();
        req.model = model;
        req.temperature = 0.6;
        req.tools = Collections.singletonList(Tool.builtinWebSearch());
        req.messages = Arrays.asList(
                new Message("system", "你是 Kimi，你每次都使用联网插件回答问题"),
                new Message("user", query)
        );
        return req;
    }

    private ChatRequest buildSecondRequest(String query, Message assistantWithToolCalls, ToolCall toolCall) {
        ChatRequest req = new ChatRequest();
        req.model = model;
        req.temperature = 0.6;
        req.tools = Collections.singletonList(Tool.builtinWebSearch());

        Message assistant = new Message();
        assistant.role = "assistant";
        assistant.content = "";
        assistant.tool_calls = Collections.singletonList(toolCall);

        Message tool = new Message();
        tool.role = "tool";
        tool.tool_call_id = toolCall.id;
        tool.name = toolCall.function != null ? toolCall.function.name : "$web_search";
        tool.content = toolCall.function != null ? toolCall.function.arguments : "{}";

        req.messages = Arrays.asList(
                new Message("system", "你是 Kimi，你每次都使用联网插件回答问题"),
                new Message("user", query),
                assistant,
                tool
        );
        return req;
    }

    private ChatResponse callKimi(ChatRequest request) {
        if (apiKey == null || apiKey.trim().isEmpty()) {
            throw new RuntimeException("KIMI API 密钥未配置或为空");
        }

        HttpHeaders headers = new HttpHeaders();
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(apiKey);
        headers.set("User-Agent", "Kimi-AI-Client/1.0");

        HttpEntity<ChatRequest> entity = new HttpEntity<>(request, headers);

        long start = System.currentTimeMillis();
        try {
            ResponseEntity<ChatResponse> resp = restTemplate.exchange(baseUrl, HttpMethod.POST, entity, ChatResponse.class);
            long cost = System.currentTimeMillis() - start;
            aiModelLogger.info("KIMI API 调用成功 - 耗时: {}ms, 状态码: {}", cost, resp.getStatusCode());
            return resp.getBody();
        } catch (Exception e) {
            long cost = System.currentTimeMillis() - start;
            aiModelLogger.error("KIMI API 调用失败 - 耗时: {}ms, 错误: {}", cost, e.getMessage());
            throw e;
        }
    }

    private long calculateBackoffDelay(int attempt) {
        long delay = baseDelayMs * (1L << (attempt - 1));
        delay += random.nextInt(1000);
        return Math.min(delay, maxDelayMs);
    }

    private AiCallResponse buildAiResponse(String content) {
        AiCallResponse res = new AiCallResponse();
        if (content == null || content.trim().isEmpty()) {
            res.setText("抱歉，KIMI 服务暂时不可用，请稍后重试。");
            res.setContents(Collections.singletonList("服务暂时不可用"));
            return res;
        }
        res.setText(content);
        res.setContents(ParseTextUtils.parseDeepseekText(content));
        return res;
    }

    // ==================== DTOs ====================

    public static class ChatRequest {
        public String model;
        public List<Message> messages;
        public Double temperature;
        public List<Tool> tools;
        public Boolean stream;
    }

    public static class Message {
        public String role;
        public String content;
        public String name;
        public String tool_call_id;
        public List<ToolCall> tool_calls;

        public Message() {}

        public Message(String role, String content) {
            this.role = role;
            this.content = content;
        }
    }

    public static class ToolCall {
        public String id;
        public String type; // builtin_function
        public ToolFunction function;
    }

    public static class ToolFunction {
        public String name; // $web_search
        public String arguments; // 原样字符串
    }

    public static class Tool {
        public String type; // builtin_function
        public ToolFunction function;

        public static Tool builtinWebSearch() {
            Tool t = new Tool();
            t.type = "builtin_function";
            ToolFunction f = new ToolFunction();
            f.name = "$web_search";
            t.function = f;
            return t;
        }
    }

    public static class ChatResponse {
        public String id;
        public String object;
        public long created;
        public String model;
        public List<Choice> choices;
        public Usage usage;
    }

    public static class Choice {
        public int index;
        public Message message;
        public String finish_reason;
    }

    public static class Usage {
        public int prompt_tokens;
        public int completion_tokens;
        public int total_tokens;
    }
}
